{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Exercise 11 - Pass Neural Net Weights Between Processes\n",
    "\n",
    "**GOAL:** The goal of this exercise is to show how to send neural network weights between workers and the driver.\n",
    "\n",
    "For more details on using Ray with TensorFlow, see the documentation at http://ray.readthedocs.io/en/latest/using-ray-with-tensorflow.html.\n",
    "\n",
    "### Concepts for this Exercise - Getting and Setting Neural Net Weights\n",
    "\n",
    "Since pickling and unpickling a TensorFlow graph can be inefficient or may not work at all, it is most efficient to ship the weights between processes as a dictionary of numpy arrays (or as a flattened numpy array).\n",
    "\n",
    "We provide the helper class `ray.experimental.TensorFlowVariables` to help with getting and setting weights. Similar techniques should work other neural net libraries.\n",
    "\n",
    "Consider the following neural net definition.\n",
    "\n",
    "```python\n",
    "import tensorflow as tf\n",
    "\n",
    "x_data = tf.placeholder(tf.float32, shape=[100])\n",
    "y_data = tf.placeholder(tf.float32, shape=[100])\n",
    "\n",
    "w = tf.Variable(tf.random_uniform([1], -1.0, 1.0))\n",
    "b = tf.Variable(tf.zeros([1]))\n",
    "y = w * x_data + b\n",
    "\n",
    "loss = tf.reduce_mean(tf.square(y - y_data))\n",
    "optimizer = tf.train.GradientDescentOptimizer(0.5)\n",
    "grads = optimizer.compute_gradients(loss)\n",
    "train = optimizer.apply_gradients(grads)\n",
    "\n",
    "init = tf.global_variables_initializer()\n",
    "sess = tf.Session()\n",
    "sess.run(init)\n",
    "```\n",
    "\n",
    "Then we can use the helper class as follows.\n",
    "\n",
    "```python\n",
    "variables = ray.experimental.TensorFlowVariables(loss, sess)\n",
    "# Here 'weights' is a dictionary mapping variable names to the associated\n",
    "# weights as a numpy array.\n",
    "weights = variables.get_weights()\n",
    "variables.set_weights(weights)\n",
    "```\n",
    "\n",
    "Note that there are analogous methods `variables.get_flat` and `variables.set_flat`, which concatenate the weights as a single array instead of a dictionary.\n",
    "\n",
    "```python\n",
    "# Here 'weights' is a numpy array of all of the neural net weights\n",
    "# concatenated together.\n",
    "weights = variables.get_flat()\n",
    "variables.set_flat(weights)\n",
    "```\n",
    "\n",
    "In this exercise, we will use an actor containing a neural network and implement methods to extract and set the neural net weights.\n",
    "\n",
    "**WARNING:** This exercise is more complex than previous exercises."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import absolute_import\n",
    "from __future__ import division\n",
    "from __future__ import print_function\n",
    "\n",
    "import numpy as np\n",
    "import ray\n",
    "import tensorflow as tf\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ray.init(num_cpus=4, include_webui=False, ignore_reinit_error=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The code below defines a class containing a simple neural network.\n",
    "\n",
    "**EXERCISE:** Implement the `set_weights` and `get_weights` methods. This should be done using the `ray.experimental.TensorFlowVariables` helper class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@ray.remote\n",
    "class SimpleModel(object):\n",
    "    def __init__(self):\n",
    "        x_data = tf.placeholder(tf.float32, shape=[100])\n",
    "        y_data = tf.placeholder(tf.float32, shape=[100])\n",
    "\n",
    "        w = tf.Variable(tf.random_uniform([1], -1.0, 1.0))\n",
    "        b = tf.Variable(tf.zeros([1]))\n",
    "        y = w * x_data + b\n",
    "\n",
    "        self.loss = tf.reduce_mean(tf.square(y - y_data))\n",
    "        optimizer = tf.train.GradientDescentOptimizer(0.5)\n",
    "        grads = optimizer.compute_gradients(self.loss)\n",
    "        self.train = optimizer.apply_gradients(grads)\n",
    "\n",
    "        init = tf.global_variables_initializer()\n",
    "        self.sess = tf.Session()\n",
    "\n",
    "        # Here we create the TensorFlowVariables object to assist with getting\n",
    "        # and setting weights.\n",
    "        self.variables = ray.experimental.TensorFlowVariables(self.loss, self.sess)\n",
    "\n",
    "        self.sess.run(init)\n",
    "\n",
    "    def set_weights(self, weights):\n",
    "        \"\"\"Set the neural net weights.\n",
    "        \n",
    "        This method should assign the given weights to the neural net.\n",
    "        \n",
    "        Args:\n",
    "            weights: Either a dict mapping strings (the variable names) to numpy\n",
    "                arrays or a single flattened numpy array containing all of the\n",
    "                concatenated weights.\n",
    "        \"\"\"\n",
    "        # EXERCISE: You will want to use self.variables here.\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def get_weights(self):\n",
    "        \"\"\"Get the neural net weights.\n",
    "        \n",
    "        This method should return the current neural net weights.\n",
    "        \n",
    "        Returns:\n",
    "            Either a dict mapping strings (the variable names) to numpy arrays or\n",
    "                a single flattened numpy array containing all of the concatenated\n",
    "                weights.\n",
    "        \"\"\"\n",
    "        # EXERCISE: You will want to use self.variables here.\n",
    "        raise NotImplementedError"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a few actors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "actors = [SimpleModel.remote() for _ in range(4)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**EXERCISE:** Get the neural net weights from all of the actors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "raise Exception('Implement this.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**EXERCISE:** Average all of the neural net weights.\n",
    "\n",
    "**NOTE:** This will be easier to do if you chose to use `get_flat`/`set_flat` instead of `get_weights`/`set_weights` in the implementation of `SimpleModel.set_weights` and `SimpleModel.get_weights` above.."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "raise Exception('Implement this.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**EXERCISE:** Set the average weights on the actors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "raise Exception('Implement this.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**VERIFY:** Check that all of the actors have the same weights."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "weights = ray.get([actor.get_weights.remote() for actor in actors])\n",
    "\n",
    "for i in range(len(weights)):\n",
    "    np.testing.assert_equal(weights[i], weights[0])\n",
    "\n",
    "print('Success! The test passed.')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.1"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": false,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
